import sys
sys.path.insert(0,"../../")

import argparse
import os

import cv2
import numpy as np

import onnxruntime

from yolox.data.data_augment import preproc as preprocess
from yolox.data.datasets import COCO_CLASSES
from yolox.utils import mkdir, multiclass_nms, demo_postprocess, vis

input_shape = (640,640)
model = "yolox_s.onnx"
score_thr = 0.3

session = onnxruntime.InferenceSession(model)

def onxx_process(origin_img):
    img, ratio = preprocess(origin_img, input_shape)

    ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
    output = session.run(None, ort_inputs)
    predictions = demo_postprocess(output[0], input_shape)[0]

    boxes = predictions[:, :4]
    scores = predictions[:, 4:5] * predictions[:, 5:]

    boxes_xyxy = np.ones_like(boxes)
    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2.
    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2.
    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2.
    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2.
    boxes_xyxy /= ratio
    dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
    if dets is not None:
        final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5]
        origin_img = vis(origin_img, final_boxes, final_scores, final_cls_inds,
                         conf=score_thr, class_names=COCO_CLASSES)

if __name__ == '__main__':
    # 创建VideoCapture对象，参数为0表示使用本地摄像头
    cap = cv2.VideoCapture(0)
    cap.set(3, 640) #width
    cap.set(4, 640) #height
    print("start capture")

    idx = 0
    while True:
        ret, frame = cap.read() #从摄像头中读取一帧图像
        idx = idx + 1
        if idx > 10:
            print(idx)
            onxx_process(frame)
        cv2.imshow('Local Camera', frame) #显示图像

        # 按下q键退出程序
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放资源
    print("capture release")
    cap.release()
    cv2.destroyAllWindows()
